"""
dataset.py
自定义 Dataset：告诉 PyTorch 如何读取一张图片并返回「预处理后的张量 + 标签」
"""
import os
from PIL import Image
from torch.utils.data import Dataset    # PyTorch 数据集基类
from torchvision import transforms      # 常用图像变换

class HerbDataset(Dataset):
    """
    构造函数
    root: 数据根目录，如 data/train 或 data/test
    input_size: 图片统一缩放尺寸
    mode: 'train' 或 'test'，决定是否需要数据增强
    """
    def __init__(self, root, input_size=224, mode='train'):
        self.paths = []   # 所有图片完整路径
        self.labels = []  # 对应标签索引
        # 类别名称按字母排序，保证顺序固定
        self.classes = sorted(os.listdir(root))
        # 类别名 -> 数字索引 的映射
        self.class_to_idx = {c: i for i, c in enumerate(self.classes)}

        # 遍历每个类别
        for cls in self.classes:
            cls_dir = os.path.join(root, cls)
            for file_name in os.listdir(cls_dir):
                # 记录图片路径
                self.paths.append(os.path.join(cls_dir, file_name))
                # 记录标签索引
                self.labels.append(self.class_to_idx[cls])

        # ---------- 数据预处理 ----------
        if mode == 'train':
            # 训练时：缩放 + 随机水平翻转（数据增强，提高泛化）
            self.trans = transforms.Compose([
                transforms.Resize((input_size, input_size)),
                transforms.RandomHorizontalFlip(p=0.5),
                transforms.ToTensor(),                      # 转成 0~1 的 Tensor
                transforms.Normalize(                       # ImageNet 标准均值方差
                    mean=[0.485, 0.456, 0.406],
                    std =[0.229, 0.224, 0.225])
            ])
        else:
            # 测试时：只做缩放和标准化，不做增强
            self.trans = transforms.Compose([
                transforms.Resize((input_size, input_size)),
                transforms.ToTensor(),
                transforms.Normalize(
                    mean=[0.485, 0.456, 0.406],
                    std =[0.229, 0.224, 0.225])
            ])

    # 必须实现：数据集大小
    def __len__(self):
        return len(self.paths)

    # 必须实现：读取第 idx 张图片
    def __getitem__(self, idx):
        img_path = self.paths[idx]
        # 用 Pillow 打开图片并转为 RGB（防止灰度图等问题）
        img = Image.open(img_path).convert('RGB')
        # 应用预处理
        img_tensor = self.trans(img)
        # 获取对应标签
        label = self.labels[idx]
        return img_tensor, label